import os
import numpy as np
from PIL import Image, ImageOps
from pathlib import Path

# ---------- 配置参数 ----------
INPUT_DIR = "D:/Desk/0312-炳哥尺寸转换/尺寸标注/标注信息"       # 输入图片目录
OUTPUT_DIR = "D:/Desk/0312-炳哥尺寸转换/尺寸标注/标注信息3"      # 输出图片目录
BLACK_THRESHOLD = 30               # 判定黑色的RGB阈值 (0~255)
MAIN_COLOR_THRESHOLD = 0.5         # 主色占比阈值 (超过此比例则判定为黑色主色)
# -----------------------------

def is_mainly_black(image, black_threshold=BLACK_THRESHOLD, ratio_threshold=MAIN_COLOR_THRESHOLD):
    """判断图片主色是否为黑色"""
    # 将图像转为NumPy数组
    img_array = np.array(image)
    # 计算黑色像素数量（RGB值均小于阈值）
    black_pixels = np.all(img_array[:, :, :3] <= black_threshold, axis=2)
    black_ratio = np.mean(black_pixels)
    return black_ratio > ratio_threshold

def invert_image(image):
    """反相颜色（跳过Alpha通道）"""
    if image.mode == 'RGBA':
        # 分离Alpha通道
        r, g, b, a = image.split()
        rgb = Image.merge('RGB', (r, g, b))
        inverted = ImageOps.invert(rgb)
        # 合并回Alpha通道
        inverted.putalpha(a)
        return inverted
    else:
        return ImageOps.invert(image)

def process_images():
    # 创建输出目录
    Path(OUTPUT_DIR).mkdir(exist_ok=True)

    # 遍历输入目录
    for filename in os.listdir(INPUT_DIR):
        input_path = os.path.join(INPUT_DIR, filename)
        output_path = os.path.join(OUTPUT_DIR, filename)

        try:
            # 打开图片
            with Image.open(input_path) as img:
                # 转换为RGB/RGBA模式
                img = img.convert('RGBA' if img.mode == 'P' else img.mode)
                # 判断主色是否为黑色
                if is_mainly_black(img):
                    # 反相处理并保存
                    inverted_img = invert_image(img)
                    inverted_img.save(output_path)
                    print(f"反相处理: {filename}")
                else:
                    # 直接复制非黑色主色图片
                    img.save(output_path)
                    print(f"跳过处理: {filename}")

        except Exception as e:
            print(f"错误处理 {filename}: {str(e)}")

if __name__ == "__main__":
    process_images()